import numpy as np
import math
import pickle
import matplotlib.pyplot as plt
from collections import deque

# 从文件加载建筑结构数据
with open("data/data_floors_0427.pkl", "rb") as haddle:
    data_floors = pickle.load(haddle)

num_floors, floor_width, floor_length, cell_size, people = data_floors[0:5] + (data_floors[11],)

# -----------------------------
# 全局参数配置
# -----------------------------
SOCIAL_FORCE_WEIGHT = 0.6      # 群体行为权重 [0-1]
CROWD_DENSITY_THRESHOLD = 3.0  # 拥挤度阈值
DISTANCE_WEIGHT = 0.7          # 距离出口权重
SAFE_DISTANCE = 2.0            # 安全距离系数
DYNAMIC_UPDATE_INTERVAL = 5    # 动态更新间隔（秒）

# -----------------------------
# 动态势能场计算（融合C1-C3）
# -----------------------------
class DynamicPotentialField:
    def __init__(self):
        self.exit_attractions = {}  # 出口吸引力场
        self.crowd_density = np.zeros((num_floors, floor_width, floor_length))  # 人群密度场
        
    def update_density(self, people):
        """实时更新人群密度场（C2）"""
        self.crowd_density.fill(0)
        for p in people:
            f, (i,j), _ = p
            self.crowd_density[f, i, j] += 1
            
    def calculate_exit_attractions(self, exits):
        """计算出口吸引力（C1）"""
        self.exit_attractions = {}
        for exit_pos in exits:
            z, x, y = exit_pos
            # 增加重力势能因子（距离越近吸引力越大）
            attraction = 1.0 / (math.hypot(x - floor_width/2, y - floor_length/2) + 1e-6)
            self.exit_attractions[exit_pos] = attraction

# -----------------------------
# 改进的疏散算法（C1-C3）
# -----------------------------
class EvacuationSimulator:
    def __init__(self):
        self.potential_field = DynamicPotentialField()
        self.social_forces = {}
        self.exit_targets = {}
        
    def assign_targets(self, people, exits):
        """动态分配目标出口（C1+C3）"""
        exit_set = set(exits)
        exit_counts = {exit:0 for exit in exit_set}
        
        # 初始随机分配出口
        for p in people:
            f, pos, _ = p
            if f == 0:  # 仅地面层人员分配出口
                exit = random.choice(list(exit_set))
                self.exit_targets[p] = exit
                exit_counts[exit] += 1
                
        # 动态调整目标出口（社会力影响）
        for p in people:
            f, pos, _ = p
            if f != 0: continue
                
            # 计算邻居平均方向（C3）
            neighbors = self.get_neighbors(pos, f)
            avg_force = np.mean([(nj-i, ni-j) for nj, ni in neighbors], axis=0) if neighbors else None
            
            # 选择吸引力最大的出口
            max_attraction = -1
            best_exit = None
            for exit_pos in exit_set:
                # 综合距离和社会力方向
                attraction = self.potential_field.exit_attractions[exit_pos]
                if avg_force is not None:
                    dir_match = np.dot(avg_force, (exit_pos[1]-pos[1], exit_pos[2]-pos[2]))
                    attraction *= (1 + dir_match)
                if attraction > max_attraction:
                    max_attraction = attraction
                    best_exit = exit_pos
                    
            self.exit_targets[p] = best_exit

    def get_neighbors(self, pos, f):
        """获取邻近个体（C3）"""
        i,j = pos
        neighbors = []
        for di in [-1,0,1]:
            for dj in [-1,0,1]:
                if di == 0 and dj == 0: continue
                ni, nj = i+di, j+dj
                if 0<=ni<floor_width and 0<=nj<floor_length:
                    for p in people:
                        if p[0] == f and np.linalg.norm(np.array(p[1])-np.array((ni,nj))) < 1.5:
                            neighbors.append(p[1])
        return neighbors

    def move_people(self, people):
        """动态路径规划（C1-C3）"""
        new_people = []
        self.potential_field.update_density(people)
        self.potential_field.calculate_exit_attractions(exits)
        
        for p in people:
            f, (i,j), _ = p
            if floors[f,i,j] == 3 and f == 0:
                continue  # 到达出口
                
            # 计算综合势能场
            potential = self.calculate_potential(f, i, j)
            
            # C3: 社会力影响
            neighbors = self.get_neighbors((i,j), f)
            social_force = self.calculate_social_force(neighbors) if neighbors else None
            
            # 计算移动方向
            move_vector = self.calculate_move_vector(potential, social_force)
            
            # 动态速度调整（C2）
            density = self.potential_field.crowd_density[f, i, j]
            speed = self.calculate_dynamic_speed(density)
            
            # 更新位置
            new_i = max(0, min(floor_width-1, i + int(move_vector[0]*speed)))
            new_j = max(0, min(floor_length-1, j + int(move_vector[1]*speed)))
            
            new_people.append([f, (new_i, new_j), speed])
            
        return new_people

    def calculate_potential(self, f, i, j):
        """综合势能计算（C1+C2）"""
        potential = 0
        # C1: 出口吸引力
        for exit_pos, strength in self.potential_field.exit_attractions.items():
            z, x, y = exit_pos
            if z == f:
                distance = math.hypot(x-i, y-j)
                potential += strength * (1 / (distance + 1e-6))
        # C2: 拥挤度惩罚
        current_density = self.potential_field.crowd_density[f, i, j]
        potential += CROWD_DENSITY_THRESHOLD / (current_density + 1e-6)
        return potential

    def calculate_social_force(self, neighbors):
        """社会力计算（C3）"""
        if not neighbors: return None
        # 计算平均方向和距离
        vectors = [np.array(n) - np.array(pos) for pos in neighbors]
        distances = [math.hypot(v[0], v[1]) for v in vectors]
        avg_vector = np.mean(vectors, axis=0)
        return avg_vector / (np.mean(distances) + 1e-6)

    def calculate_move_vector(self, potential, social_force):
        """移动方向计算"""
        # 势能梯度计算
        grad_potential = np.gradient(potential)
        # 社会力方向
        social_dir = social_force / np.linalg.norm(social_force) if social_force is not None else None
        # 方向合成
        move_dir = grad_potential + (social_dir * SOCIAL_FORCE_WEIGHT if social_dir is not None else 0)
        return move_dir / np.linalg.norm(move_dir) if np.linalg.norm(move_dir) > 0 else np.random.rand(2)

    def calculate_dynamic_speed(self, density):
        """动态速度调整（C2）"""
        base_speed = 1.2  # 基准速度
        if density < 2:
            return base_speed * 1.2
        elif density < 4:
            return base_speed * 0.8
        else:
            return base_speed * 0.5

# -----------------------------
# 主仿真循环
# -----------------------------
simulator = EvacuationSimulator()
exits = [(0, pos) for pos in data_floors[10]]

# 初始化目标分配
simulator.assign_targets(people, exits)

# 创建动画
fig, axes = plt.subplots(1, num_floors, figsize=(5*num_floors,5))
time_elapsed = 0

def update(frame):
    global time_elapsed
    time_elapsed += DYNAMIC_UPDATE_INTERVAL
    
    # 动态路径规划
    people = simulator.move_people(people)
    
    # 绘制建筑结构
    for f in range(num_floors):
        ax = axes[f]
        ax.imshow(floors[f]==1, cmap='Greys', vmin=0, vmax=1)
        
        # 绘制出口和楼梯
        for exit_pos in exits:
            if exit_pos[0] == f:
                ax.plot(exit_pos[1], exit_pos[2], 'rs', markersize=10)
        for stair in stairs_up[f] + stairs_down[f]:
            ax.plot(stair[1], stair[0], 'gv' if stair in stairs_up[f] else 'bv', markersize=8)
    
    # 绘制人员分布
    for p in people:
        f, (i,j), _ = p
        if f == 0 and floors[f,i,j] == 3: continue
        axes[f].scatter(j, i, c='r', s=10)
    
    # 更新标题
    plt.suptitle(f'Evacuation Progress - Time: {time_elapsed:.1f}s')
    return fig,

# 创建动画
ani = animation.FuncAnimation(fig, update, frames=range(120), interval=500, blit=False)

plt.show()